package com.boboPlanet.comm.thread;

import com.boboPlanet.comm.context.TraceIdContext;
import com.boboPlanet.comm.enums.PlanetCommCode;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.*;

/**
 * @Description: TODO 解决feign调用时MDC获取不到问题、异步traceId传递问题
 * @Author: mis_wu
 * @Date: 2023/5/13
 * @Company planet
 **/
public class ThreadMdcUtil {

    public static void setTraceIdIfAbsent() {
        if (get() == null || StringUtils.isBlank(get())) {
            String traceId = TraceIdAutoGenerate.traceByUuid();
            MDC.put(PlanetCommCode.COMM_LOG_KEY.getCode(), traceId);
            TraceIdContext.setTraceId(traceId);
        }
    }

    public static String setTraceId() {
        String traceId = TraceIdAutoGenerate.traceByUuid();
        MDC.put(PlanetCommCode.COMM_LOG_KEY.getCode(), traceId);
        TraceIdContext.setTraceId(traceId);

        return traceId;
    }

    public static void setTraceId(String traceId) {
        MDC.put(PlanetCommCode.COMM_LOG_KEY.getCode(), traceId);
        TraceIdContext.setTraceId(traceId);
    }

    public static String get() {
        return MDC.get(PlanetCommCode.COMM_LOG_KEY.getCode());
    }

    public static void remove() {
        if (get() != null) {
            MDC.remove(PlanetCommCode.COMM_LOG_KEY.getCode());
            TraceIdContext.remove();
        }
    }

    public static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
        return () -> {
            if (context == null) {
                MDC.clear();
                TraceIdContext.remove();
            } else {
                MDC.setContextMap(context);
            }
            setTraceIdIfAbsent();
            try {
                return callable.call();
            } finally {
                MDC.clear();
                TraceIdContext.remove();
            }
        };
    }

    public static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
        return () -> {
            if (context == null) {
                MDC.clear();
                TraceIdContext.remove();
            } else {
                MDC.setContextMap(context);
            }
            setTraceIdIfAbsent();
            try {
                runnable.run();
            } finally {
                MDC.clear();
                TraceIdContext.remove();
            }
        };
    }

}
